# 下载股票数据
import akshare as ak
import pandas as pd
import mplfinance as mpf
import matplotlib.pyplot as plt
import numpy as np

# A 股日频率数据-东方财富 stock_zh_a_hist
'''
    '日期': 'Date',
    '开盘': 'Open',
    '最高': 'High',
    '最低': 'Low',
    '收盘': 'Close',
    '成交量': 'Volume',
    '成交额':
    '振幅':
    '涨跌幅': 
    '涨跌额':
    '换手率':
'''

stock = ak.stock_zh_a_hist(symbol="002236",
                           period="daily",
                           start_date="20220501",
                           end_date="20221231", adjust="hfq")

stock = pd.DataFrame(stock, columns=['日期', '开盘', '最高', '最低', '收盘', '成交量'])
stock.rename(columns={
    '日期': 'date',
    '开盘': 'open',
    '最高': 'high',
    '最低': 'low',
    '收盘': 'close',
    '成交量': 'volume'
},
    inplace=True)
stock.index = pd.DatetimeIndex(stock['date'])
print('----原数据------')
print(stock.head(5))
'''
# 绘制蜡烛图
# type='candle', type='line', type='renko', or type='pnf'
'''
mpf.plot(stock.tail(30), type="candle", volume=True)

'''
简单数据处理
- 相比前一个交易日收益计算
'''
stock['diff'] = stock['close'].diff()
print('----计算涨跌diff数据------')
print(stock.head(5))

'''
简单交易策略
·当日股价下跌，下一个交易日买入
·当日股价上涨，下一个交易日卖出

创建交易信号字段：Signal, diff > 0 Signal=1 卖出，否则Signal=0
'''

stock['signal'] = np.where(stock['diff'] > 0, 1, 0)
print('----计算买入卖出信号后的数据------')
print(stock.head(5))

# matplotlib 绘图
plt.figure(figsize=(10, 5))
stock['close'].plot(linewidth=2, color='k', grid=True)
# 卖出标志 x轴日期，y轴数值
# matplotlib.pyplot.scatter(x, y, marker, size, color)
plt.scatter(stock['close'].loc[stock.signal == 1].index,
            stock['close'][stock.signal == 1],
            marker='v', s=80, c='g')
# 买入标志
plt.scatter(stock['close'].loc[stock.signal == 0].index,
            stock['close'][stock.signal == 0],
            marker='^', s=80, c='r')
plt.show()
